package com.aikg.kgrag_java.service.adapter;

import com.aikg.kgrag_java.pojo.dto.neo4j.Neo4jRetrievalDTO;
import com.aikg.kgrag_java.pojo.entity.neo4j.Node;
import com.aikg.kgrag_java.pojo.vo.Neo4jRetrievalVO;
import com.aikg.kgrag_java.pojo.vo.NodeVO;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @ClassName: Neo4jRetrievalAdapter
 * @Description: 适配器将dto转为vo
 * @Author: 李宗雨
 * @Date: 2025/1/6 14:36
 * @Version: 1.0
 **/
public class Neo4jRetrievalAdapter {

    public static Neo4jRetrievalVO build(Neo4jRetrievalDTO neo4jRetrievalDTO) {
        Neo4jRetrievalVO vo = new Neo4jRetrievalVO();

        for (Node node : neo4jRetrievalDTO.getNodeInfoList()) {
            if (node.getName().isEmpty() && node.getProperty().isEmpty()) {
                continue;
            }
            NodeVO nodeVO = new NodeVO();
            nodeVO.setName(node.getName());
            Map<String, String> temp = new HashMap<>();
            Map<String, String> propertyMap = node.getProperty();
            for (String key : propertyMap.keySet()) {
                if (!propertyMap.get(key).isEmpty() && !propertyMap.get(key).equals("null")) {
                    temp.put(key, propertyMap.get(key));
                }
            }
            nodeVO.setAttributes(temp);
            vo.getNode_info_list().add(nodeVO);
        }

        // 使用Stream API进行转换
        vo.getRelation_triplet_list().addAll(neo4jRetrievalDTO.getRelationTripletList().stream()
                .map(map -> {
                    Map<String, String> newMap = new HashMap<>();
                    newMap.put("source", map.get("头实体"));
                    newMap.put("target", map.get("尾实体"));
                    newMap.put("value", map.get("关系"));
                    return newMap;
                })
                .collect(Collectors.toList()));
        return vo;
    }
}
